"""中性可复用的动图生成工具

提供一个不含硬编码的函数 `generate_density_gif`，用于把若干年份的 GeoDataFrame 按列绘制为若干帧并合成 GIF。

设计要点：
- 使用 `src.plot.prepare_border_counties_population_data` 做数据准备（复用已有逻辑）。
- 参数化所有可配置项（列名、色带、dpi、figsize、colorbar bbox、字体名等）。
- 可选把每帧导出为 PNG（用于 QA）。
- 使用 Pillow 的 save_all 写入多帧 GIF，遇到异常回退到 imageio。
"""
from pathlib import Path
from typing import Iterable, Optional, Tuple
import io
import numpy as np
import pandas as pd
from PIL import Image as PILImage
import imageio
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter

from src.plot import prepare_border_counties_population_data
from utils.plot_style import set_plot_style, set_chinese_font, force_chinese_font


def _format_cb_ticks_per_ten_thousand(x, pos):
    val = x / 10000.0
    if val >= 10:
        return f"{val:.0f}"
    return f"{val:.1f}"


def generate_density_gif(
    years: Iterable[int],
    out_path: str,
    population_parquet: str,
    column: str = 'mean',
    scope: str = 'county',
    cmap: str = 'Reds',
    dpi: int = 150,
    figsize: Tuple[int, int] = (10, 12),
    duration: float = 0.6,
    font_name: Optional[str] = None,
    save_frames: bool = False,
    frames_dir: str = 'figures/tmp_frames',
    verbose: bool = True,
) -> Tuple[str, int]:
    """生成动图并返回 (gif_path, frame_count)

    参数尽量中性化并可复用。
    """
    outp = Path(out_path)
    outp.parent.mkdir(parents=True, exist_ok=True)

    # 先收集每年的预处理结果并计算全局 vmin/vmax
    geos = {}
    vmin = None
    vmax = None
    for y in years:
        prep = prepare_border_counties_population_data(year=y, scope=scope, population_parquet=population_parquet)
        if not prep:
            if verbose:
                print(f"year {y}: no data, skip")
            continue
        geos[y] = prep
        geo = prep['geo']
        if column in geo.columns:
            ser = pd.to_numeric(geo[column], errors='coerce').fillna(0)
        elif 'mean' in geo.columns:
            ser = pd.to_numeric(geo['mean'], errors='coerce').fillna(0)
        else:
            ser = pd.Series(dtype=float)
        if not ser.empty:
            yr_min = float(ser.min())
            yr_max = float(ser.max())
            if vmin is None or yr_min < vmin:
                vmin = yr_min
            if vmax is None or yr_max > vmax:
                vmax = yr_max

    if not geos:
        raise RuntimeError('No data found for any year')

    if vmin is None or vmax is None:
        vmin, vmax = 0.0, 1.0
    if vmin == vmax:
        vmax = vmin + 1e-6

    frames = []
    frames_dir_path = Path(frames_dir)
    if save_frames:
        frames_dir_path.mkdir(parents=True, exist_ok=True)

    for y in sorted(geos.keys()):
        prep = geos[y]
        geo = prep['geo']
        plot_cfg = prep.get('cfg_plot') or {}
        set_plot_style(plot_cfg)

        plot_figsize = prep.get('figsize', figsize)
        plot_dpi = prep.get('dpi', dpi)

        fig, ax = plt.subplots(figsize=figsize or plot_figsize, dpi=plot_dpi)
        # 选列
        if column in geo.columns:
            col = column
        else:
            col = 'mean' if 'mean' in geo.columns else '总人口'

        geo.plot(column=col, cmap=cmap, linewidth=0.2, edgecolor='#444444', ax=ax, vmin=vmin, vmax=vmax)
        ax.set_axis_off()

        mappable = mpl.cm.ScalarMappable(norm=mpl.colors.Normalize(vmin=vmin, vmax=vmax), cmap=cmap)
        mappable._A = []
        cax = fig.add_axes(prep['colorbar_bbox'])
        cb = fig.colorbar(mappable, cax=cax, orientation='horizontal')
        cb.ax.xaxis.set_major_formatter(FuncFormatter(_format_cb_ticks_per_ten_thousand))
        cb.ax.tick_params(labelsize=9)
        if col == 'mean':
            label_text = '贫困人口密度'
        else:
            label_text = '人口密度（单位根据数据）'
        cb.set_label(label_text, fontsize=11)

        # 字体（若传入 font_name，则尝试设置）
        if font_name:
            try:
                set_chinese_font(font_name)
                force_chinese_font(ax, font_name)
                for lab in cb.ax.get_xticklabels():
                    lab.set_fontfamily(font_name)
                cb.ax.xaxis.label.set_fontfamily(font_name)
            except Exception:
                pass

        # 年份注记（左上角）
        try:
            txt_kwargs = dict(color='white', fontsize=20, weight='bold', va='top', ha='left')
            ax.text(0.01, 0.99, str(y), transform=ax.transAxes, bbox=dict(facecolor='black', alpha=0.35, pad=4), **txt_kwargs)
        except Exception:
            pass

        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=plot_dpi, bbox_inches='tight')
        plt.close(fig)
        buf.seek(0)
        pil_img = PILImage.open(buf).convert('RGBA')
        frames.append(pil_img.copy())
        buf.close()

        if save_frames:
            # save png per frame for QA
            frame_path = frames_dir_path / f"frame_{y}.png"
            pil_img.save(frame_path)

    # write gif
    try:
        if len(frames) == 1:
            frames[0].save(outp, format='GIF')
        else:
            first, rest = frames[0], frames[1:]
            first.save(outp, format='GIF', save_all=True, append_images=rest, duration=int(duration*1000), loop=0)
    except Exception:
        # fallback
        imageio.mimsave(outp, [np.array(f) for f in frames], format='GIF', duration=duration, loop=0)

    if verbose:
        print(f"Saved GIF to {outp} (frames={len(frames)})")
    return str(outp), len(frames)
